# -*- coding: utf-8 -*-

import argparse
import subprocess
import math
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np

import float_tensor_pb2


class CCustomNet(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(CCustomNet, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True))
        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.uniform_() * 10
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.uniform_() * 10
                m.bias.data.uniform_() * 10
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

    def forward(self, x):
        x = self.conv(x)
        return x


def save_float_tensor(float_tensor_np, save_file_path):
    assert isinstance(float_tensor_np, np.ndarray)
    assert float_tensor_np.ndim == 3

    float_tensor_pb = float_tensor_pb2.FloatTensor()
    float_tensor_pb.shape.channel = float_tensor_np.shape[0]
    float_tensor_pb.shape.height = float_tensor_np.shape[1]
    float_tensor_pb.shape.width = float_tensor_np.shape[2]
    float_tensor_pb.data.extend(float_tensor_np.flatten())

    with open(save_file_path, 'wb') as save_file:
        save_file.write(float_tensor_pb.SerializeToString())


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--in_channels", type=int)
    parser.add_argument("--out_channels", type=int)
    parser.add_argument("--kernel_size", type=int)
    parser.add_argument("--stride", type=int)
    parser.add_argument("--padding", type=int)

    args = parser.parse_args()
    in_channels = args.in_channels
    out_channels = args.out_channels
    kernel_size = args.kernel_size
    stride = args.stride
    padding = args.padding

    model = CCustomNet(in_channels=in_channels,
                       out_channels=out_channels,
                       kernel_size=kernel_size,
                       stride=stride,
                       padding=padding)
    model.eval()

    input_x = Variable(torch.randn(1, in_channels, 288, 320) * 10 + 1)
    output_y = model(input_x)

    print('output 10 front values: ')
    print(output_y.data.numpy()[0][0, 0, 0:10])

    input_file_path = '../../test_data/convolution_test_case_input.pb'
    output_file_path = '../../test_data/convolution_test_case_output.pb'
    onnx_file_path = '../../test_data/convolution_test_case.onnx'
    save_float_tensor(input_x.data.numpy()[0], input_file_path)
    save_float_tensor(output_y.data.numpy()[0], output_file_path)
    torch.onnx.export(model, input_x, onnx_file_path, verbose=False)

    subprocess.call(['../onnx2ncnn',
                     '../../test_data/convolution_test_case.onnx',
                     '../../test_data/convolution_test_case.param',
                     '../../test_data/convolution_test_case.bin'])
    subprocess.call(['rm', '../../test_data/convolution_test_case.onnx'])


if __name__ == "__main__":
    main()
